from sklearn.metrics import mean_squared_error, mean_absolute_error
import numpy as np
from common.utils.mlmodel_util import get_model_info_sklearn
from common.utils.model_util import load_model
from common.utils.mlmodel_util import preprocess
from sklearn import svm, model_selection
from common.utils.format_util import dup_name_handler


def predict(dataframe, para_list, output):
    # --------------------load the model-----------------
    model_save_path, _, feature_list = get_model_info_sklearn(para_list["model_id"])
    mlmodel = load_model(model_save_path)
    # --------------------make predictions----------------
    X = preprocess(dataframe, para_list, feature_list)
    pred = mlmodel.predict(X)
    output_cols = "_prediction_"
    output_cols = dup_name_handler(output_cols, para_list["feature_col"] + [para_list["label_col"]])
    output["result"]["output_params"]["output_cols"] = output_cols
    dataframe[output_cols] = pred
    return dataframe


def train(dataframe, para_list, record):
    # --------------------prepare data------------------
    feature_list = []
    X = preprocess(dataframe, para_list, feature_list)
    label_col = para_list["label_col"]
    y = dataframe[label_col].values
    X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y,
                                                                        test_size=para_list["split_rate"][1],
                                                                        random_state=0)

    # ---------------------model fit---------------------
    mlmodel = svm.SVR(kernel=para_list["kernel"]).fit(X_train, y_train)
    y_pred = mlmodel.predict(X_test)

    # ---------------------get metrics--------------------
    n = len(y_test)
    p = len(feature_list)

    MSE = mean_squared_error(y_test, y_pred)
    RMSE = np.sqrt(MSE)
    MAE = mean_absolute_error(y_test, y_pred)
    r2 = mlmodel.score(X, y_test)
    r2adj = 1 - ((1 - r2) * (n - 1)) / (n - p - 1)

    # ---------------------record info----------------------
    other_info = {
        "feature_list": feature_list,
        "feature_col": para_list["feature_col"],
        "coefficients": list(mlmodel.coef_),
        "intercept": mlmodel.intercept_,
        "MSE": MSE,
        "MAE": MAE,
        "RMSE": RMSE,
        "r2": r2,
        "r2adj": r2adj,
        "total_iterations": 1,
        "df": n - p - 1,
    }
    record["other_info"] = other_info

    output_cols = "_prediction_"
    # output["result"]["output_params"]["output_cols"] = output_cols
    dataframe[output_cols] = y_pred
    return mlmodel, dataframe
